﻿#include <QApplication>
#include <QTableView>
#include <QHeaderView>
#include <QStringList>
#include <QVector>
#include <QStandardItemModel>
#include <QSortFilterProxyModel>
#include <QMenu>
#include <QAction>
#include <QFileDialog>
#include <QFile>
#include <QTextStream>
#include <QContextMenuEvent>

class MyTableView : public QTableView {
    Q_OBJECT

public:
    MyTableView(QWidget* parent = nullptr) : QTableView(parent) {
        setSortingEnabled(true);
        // Connect the sectionDoubleClicked signal to the sortTable slot
        connect(horizontalHeader(), &QHeaderView::sectionDoubleClicked, this, &MyTableView::sortTable);
    }

    void setTableData(const QVector<QVector<double>>& columnData, const QStringList& columnNames) {
        columenames = columnNames;
        int rows = columnData.isEmpty() ? 0 : columnData[0].size();
        int cols = columnData.size();

        QStandardItemModel* model = new QStandardItemModel(rows, cols, this);
        model->setHorizontalHeaderLabels(columnNames);

        for (int col = 0; col < cols; ++col) {
            for (int row = 0; row < rows; ++row) {
                QStandardItem* item = new QStandardItem();
                item->setData(columnData[col][row], Qt::DisplayRole);
                model->setItem(row, col, item);
            }
        }

        QSortFilterProxyModel* proxyModel = new QSortFilterProxyModel(this);
        proxyModel->setSourceModel(model);
        proxyModel->setSortRole(Qt::DisplayRole); // Ensure sorting uses the correct data role
        setModel(proxyModel);
    }

    QVector<QVector<double>> getSortedData() const {
        QVector<QVector<double>> sortedData;
        QSortFilterProxyModel* proxyModel = qobject_cast<QSortFilterProxyModel*>(model());
        if (proxyModel) {
            QStandardItemModel* model = qobject_cast<QStandardItemModel*>(proxyModel->sourceModel());
            if (model) {
                int rows = proxyModel->rowCount();
                int cols = proxyModel->columnCount();
                sortedData.resize(cols);

                for (int col = 0; col < cols; ++col) {
                    sortedData[col].resize(rows);
                    for (int row = 0; row < rows; ++row) {
                        QModelIndex proxyIndex = proxyModel->index(row, col);
                        QModelIndex sourceIndex = proxyModel->mapToSource(proxyIndex);
                        sortedData[col][row] = model->data(sourceIndex, Qt::DisplayRole).toDouble();
                    }
                }
            }
        }
        return sortedData;
    }
    QStringList getColumnNames() const {
        return columenames;
    }
protected:
    void contextMenuEvent(QContextMenuEvent* event) override {
        QMenu contextMenu(this);
        QAction saveAction("Save Sorted Data", this);
        connect(&saveAction, &QAction::triggered, this, &MyTableView::saveSortedData);
        contextMenu.addAction(&saveAction);
        contextMenu.exec(event->globalPos());
    }

private slots:
    void sortTable(int column) {
        QSortFilterProxyModel* proxyModel = qobject_cast<QSortFilterProxyModel*>(model());
        if (proxyModel) {
            Qt::SortOrder currentOrder = horizontalHeader()->sortIndicatorOrder();
            Qt::SortOrder newOrder = (currentOrder == Qt::AscendingOrder) ? Qt::DescendingOrder : Qt::AscendingOrder;
            proxyModel->sort(column, newOrder);
        }
    }
    void saveSortedData() {
        QVector<QVector<double>> sortedData = getSortedData();
        QString fileName = QFileDialog::getSaveFileName(this, "Save Sorted Data", "", "Text Files (*.txt)");
        if (!fileName.isEmpty()) {
            QFile file(fileName);
            if (file.open(QIODevice::WriteOnly | QIODevice::Text)) {
                QTextStream out(&file);
                for (auto name : columenames) {
                    out << name;
                    if (columenames.indexOf(name) < columenames.size() - 1)
                        out << "\t";
                }
                out << "\n";
                for (int row = 0; row < sortedData[0].size(); ++row) {
                    for (int col = 0; col < sortedData.size(); ++col) {
                        out << sortedData[col][row];
                        if (col < sortedData.size() - 1)
                            out << "\t";
                    }
                    out << "\n";
                }
                file.close();
            }
        }
    }

private:
    QStringList columenames;
};

